/**
 * 
 */
package com.gragra.sampling.nonproj;
import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntIterator;
import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
import com.gragra.sampling.LogarithmicFactorialPlusAlphaCash;
import com.gragra.sampling.binary.HeadAssignmentStructure;
import com.gragra.sampling.nonproj.structures.nonprojective.NonProjDepdTree;
import com.gragra.util.OpenUnitElementCounter;
import com.gragra.util.StringToInt;
/**
 * @author Christoph Teichmann
 * 23.10.2013
 */
public class TwoLevelDistHead extends DistHeadConstraint
{
	/**
	 * @param population
	 * @param tags
	 * @param tagTagAlpha
	 * @param tagDistAlpha
	 * @param fertAlpha
	 * @param nounWordConstraint
	 * @param functionWordConstraint
	 */
	public TwoLevelDistHead(int population, IntSet tags, double tagTagAlpha,
			double tagDistAlpha, double fertAlpha, double nounConstraint,
			double functionConstraint, double adjectiveConstraint,
			double determinerConstraint, double verbConstraint)
	{
		super(population, tags, tagTagAlpha, tagDistAlpha, fertAlpha,
				nounConstraint, functionConstraint, adjectiveConstraint,
				determinerConstraint, verbConstraint);
		IntIterator iit = tags.iterator();
		IntOpenHashSet seen = new IntOpenHashSet();
		while(iit.hasNext())
		{
			int tag = iit.nextInt();
			int num = Integer.parseInt(StringToInt.reverse(tag))/2;
			if(!seen.contains(num))
			{
				int n;
				tagTagSecUnits.put(StringToInt.ensure(Integer.toString(num*2)),n = getFreeUnit());
				tagTagSecUnits.put(StringToInt.ensure(Integer.toString(num*2+1)), n);
				tagDistSecUnits.put(StringToInt.ensure(Integer.toString(num*2)), n = getFreeUnit());
				tagDistSecUnits.put(StringToInt.ensure(Integer.toString(num*2+1)), n);
				fertSecUnits.put(StringToInt.ensure(Integer.toString(num*2)), n = getFreeUnit());
				fertSecUnits.put(StringToInt.ensure(Integer.toString(num*2+1)), n);
			}
			seen.add(num);
		}
		tagTagSecUnits.put(ROOTNODETAG, getFreeUnit());
		tSecUnits = new IntOpenHashSet(tagTagSecUnits.values());
		fertSecUnits.put(ROOTNODETAG, getFreeUnit());
		fSecUnits = new IntOpenHashSet(fertSecUnits.values());
		dSecUnits = new IntOpenHashSet(tagDistSecUnits.values());
		this.distSecElement = new LogarithmicFactorialPlusAlphaCash(tagDistAlpha, population);
		this.distSecUnit = new LogarithmicFactorialPlusAlphaCash(tagDistAlpha*((DISTANCES*2)+1), population);
		this.tagSecElement = new LogarithmicFactorialPlusAlphaCash(tagTagAlpha, population);
		this.tagSecUnit = new LogarithmicFactorialPlusAlphaCash(tagTagAlpha*tags.size()/2, population);
		this.fertSecUnit = new LogarithmicFactorialPlusAlphaCash(fertAlpha*((MAXCHILDREN+1)), population);
		this.fertSecElement = new LogarithmicFactorialPlusAlphaCash(fertAlpha, population);
	}
	/**
	 * 
	 */
	private static final int	DISTANCES	= 10;
	/**
	 * 
	 */
	private static final int ROOTNODETAG	= -1;
	/**
	 * 
	 */
	private static final int MAXCHILDREN	= 10;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash distSecUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash distSecElement;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash tagSecUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash tagSecElement;
	/**
	 * 
	 */
	private final IntSet tSecUnits ;
	/**
	 * 
	 */
	private final IntSet dSecUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap tagTagSecUnits = new Int2IntOpenHashMap();
	/**
	 * 
	 */
	private final Int2IntOpenHashMap tagDistSecUnits = new Int2IntOpenHashMap();
	/**
	 *
	 */
	private final Int2IntOpenHashMap fertSecUnits = new Int2IntOpenHashMap();
	/**
	 * 
	 */
	private final IntSet fSecUnits;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash	fertSecElement;
	/**
	 * 
	 */
	private LogarithmicFactorialPlusAlphaCash	fertSecUnit;
	/**
	 * 
	 * @param i
	 * @param oue
	 * @param has
	 * @param thread
	 * @param amount 
	 */
	protected void addFeatures(int pos, OpenUnitElementCounter oue,
			HeadAssignmentStructure has, int thread, double amount, IntArrayList children)
	{
		int head = has.getHeadPosition(pos, thread);
		int headTag = ROOTNODETAG;
		int dist = 0;
		if(head >= 0 && head != has.getDependencyRootPosition())
		{
			headTag = has.getTag(head);
			dist = Math.min(Math.max(head-pos,-DISTANCES),DISTANCES);
		}
		int tag = has.getTag(pos);
		int hhUnit = getSecHeadToHeadUnit(tag);
		int dUnit = getSecDistUnit(tag);
		int fUnit = getSecFertilityUnit(tag);
		oue.add(fUnit, makeFertility(pos, children),amount);
		if(pos == 0)
		{
			fUnit = getSecFertilityUnit(ROOTNODETAG);
			oue.add(fUnit,makeFertility(has.getDependencyRootPosition(), children),amount);
		}
		if(head >= 0)
		{
			oue.add(hhUnit, headTag, amount);
			oue.add(dUnit, dist, amount);
		}
		super.addFeatures(pos, oue, has, thread, amount, children);
	}
	/**
	 * 
	 * @param tag
	 * @return
	 */
	private int getSecFertilityUnit(int tag)
	{return this.fertSecUnits.get(tag);}
	/**
	 * @param tag
	 * @return
	 */
	protected int getSecHeadToHeadUnit(int tag)
	{return this.tagTagSecUnits.get(tag);}
	/**
	 * @param tag
	 * @return
	 */
	protected int getSecDistUnit(int tag)
	{return this.tagDistSecUnits.get(tag);}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getUnitCash(int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getUnitCash(int unit)
	{
		if(this.dSecUnits.contains(unit))
		{return this.distSecUnit;}
		else if (this.tSecUnits.contains(unit)) 
		{return this.tagSecUnit;}
		else if(this.fSecUnits.contains(unit))
		{return this.fertSecUnit;}
		return super.getUnitCash(unit);
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getElementCash(int, int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getElementCash(int unit,
			int element)
	{
		if(this.dSecUnits.contains(unit))
		{return this.distSecElement;}
		else if (this.tSecUnits.contains(unit)) 
		{return this.tagSecElement;}
		else if(this.fSecUnits.contains(unit))
		{return this.fertSecElement;}
		return super.getElementCash(unit, element);
	}
	@Override
	public double makeChartProposalProb(
			NonProjDepdTree tree,
			int headCandidate, int child, int thread)
	{
		OpenUnitElementCounter oue = this.getSeen(thread);
		int tag = tree.getTag(child);
		int headTag = headCandidate == tree.getRootPosition() ? ROOTNODETAG : tree.getTag(headCandidate);
		int hunit = this.getSecHeadToHeadUnit(tag);
		oue.add(hunit, headTag, 1.0);
		int dist = headCandidate == tree.getRootPosition() ? 0 : Math.min(Math.max(headCandidate-child,-DISTANCES),DISTANCES);
		int dUnit = this.getSecDistUnit(tag);
		oue.add(dUnit, dist, 1.0);
		double ret = this.makeLogProb(oue, thread);
		return ret+super.makeChartProposalProb(tree, headCandidate, child, thread);
	}
}